#include <cmath>
#include "env.hpp"
#include "tracelog.h"
constexpr double SIGN(double x) { return x>0? 1 : -1; }

PendulumEnvironment::PendulumEnvironment(double th, double dth) {
    SetLogLevel(LOG_FATAL);
    _intTheta = new simucpp::UIntegrator(&_sim1);
    _intDTheta = new simucpp::UIntegrator(&_sim1);
    _fcnThetaCycle = new simucpp::UFcn(&_sim1);
    _misoDTheta = new simucpp::UFcnMISO(&_sim1);
    _cnstPower = new simucpp::UConstant(&_sim1);
    _sim1.connectU(_misoDTheta, _intDTheta);
    _sim1.connectU(_intDTheta, _intTheta);
    _sim1.connectU(_intTheta, _fcnThetaCycle);
    _sim1.connectU(_cnstPower, _misoDTheta);
    _sim1.connectU(_fcnThetaCycle, _misoDTheta);
    _fcnThetaCycle->Set_Function([](double u){
        while (u > M_PI) u -= 2*M_PI;
        while (u < -M_PI) u += 2*M_PI;
        return u;
    });
    _misoDTheta->Set_Function([](double *u){
        return u[0] - 10 * sin(u[1]);
    });
    _sim1.Initialize();
    _sim1.Simulate_FirstStep();
}
double PendulumEnvironment::Get_Reward() {
    double h = - cos(_fcnThetaCycle->Get_OutValue());
    return h>0? 5*h : h;
}
std::vector<double> PendulumEnvironment::Get_State() {
    return std::vector<double>{_fcnThetaCycle->Get_OutValue(), _intDTheta->Get_OutValue()};
}
std::vector<double> PendulumEnvironment::Reset() {
    _sim1.Simulation_Reset();
    return Get_State();
}
std::vector<double> PendulumEnvironment::Step(double action) {
    _cnstPower->Set_OutValue(action*10-5);
    for (size_t i = 0; i < 10; i++)
        _sim1.Simulate_OneStep();
    return std::vector<double>{
        _fcnThetaCycle->Get_OutValue(),
        _intDTheta->Get_OutValue(),
        Get_Reward(), 0
    };
}
